import tensorflow as tf
# ----------------------------------------------------------------------- #
# 多周期余弦退火衰减
# ----------------------------------------------------------------------- #
# eager模式防止graph报错
# tf.config.experimental_run_functions_eagerly(True)
# ------------------------------------------------ #
import math
import numpy

# 继承自定义学习率的类
class CosineWarmupDecay(tf.keras.optimizers.schedules.LearningRateSchedule):
    '''
    initial_lr: 初始的学习率
    min_lr: 学习率的最小值
    max_lr: 学习率的最大值
    warmup_step: 线性上升部分需要的step
    total_step: 第一个余弦退火周期需要对总step
    multi: 下个周期相比于上个周期调整的倍率
    print_step: 多少个step并打印一次学习率
    '''

    # 初始化
    def __init__(self, initial_lr, min_lr, warmup_step, total_step, multi, print_step):
        # 继承父类的初始化方法
        super(CosineWarmupDecay, self).__init__()

        # 属性分配
        self.initial_lr = tf.cast(initial_lr, dtype=tf.float32)
        self.min_lr = tf.cast(min_lr, dtype=tf.float32)
        self.warmup_step = warmup_step  # 初始为第一个周期的线性段的step
        self.total_step = total_step  # 初始为第一个周期的总step
        self.multi = multi
        self.print_step = print_step

        # 保存每一个step的学习率
        self.learning_rate_list = []
        # 当前步长
        self.step = 0

    # 前向传播, 训练时传入当前step，但是上面已经定义了一个，这个step用不上
    @tf.function
    def __call__(self, step):

        # 如果当前step达到了当前周期末端就调整
        if self.step >= self.total_step:
            # 乘上倍率因子后会有小数，这里要注意
            # 调整一个周期中线性部分的step长度
            self.warmup_step = self.warmup_step * (1 + self.multi)
            # 调整一个周期的总step长度
            self.total_step = self.total_step * (1 + self.multi)

            # 重置step，从线性部分重新开始
            self.step = 0

        # 余弦部分的计算公式
        decayed_learning_rate = self.min_lr + 0.5 * (self.initial_lr - self.min_lr) * \
                                (1 + tf.math.cos(math.pi * (self.step - self.warmup_step) / (self.total_step - self.warmup_step)))

        # 计算线性上升部分的增长系数k
        k = (self.initial_lr - self.min_lr) / self.warmup_step
        # 线性增长线段 y=kx+b
        warmup = k * self.step + self.min_lr

        # 以学习率峰值点横坐标为界，左侧是线性上升，右侧是余弦下降
        decayed_learning_rate = tf.where(self.step < self.warmup_step, warmup, decayed_learning_rate)

        # 每个epoch打印一次学习率
        #if step % self.print_step == 0:
            # 打印当前step的学习率
            # print('learning_rate has changed to: ', decayed_learning_rate.numpy().item())
                # 用tf.math.equal替代Python的布尔比较
        should_print = tf.math.equal(step % self.print_step, 0)
        
        # 使用tf.cond来处理条件判断
        def print_lr():
            tf.print('learning_rate has changed to: ', decayed_learning_rate)
        
        # 仅在应该打印时执行
        tf.cond(should_print, print_lr, lambda: None)


        # 每个step保存一次学习率
        self.learning_rate_list.append(decayed_learning_rate)
        # 计算完当前学习率后step加一用于下一次
        self.step = self.step + 1

        # 返回调整后的学习率
        return decayed_learning_rate
